package com.dyface;

import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.nd4j.linalg.ops.transforms.Transforms;
import org.nd4j.linalg.schedule.ExponentialSchedule;
import org.nd4j.linalg.schedule.ScheduleType;
import org.nd4j.linalg.indexing.NDArrayIndex;

import java.io.*;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

/**
 * 人脸识别模型类
 */
public class FaceRecognitionModel {
    private MultiLayerNetwork 模型;
    private Map<Integer, String> 人脸映射表;
    private Map<Integer, List<INDArray>> 人脸特征库;
    private static final int 图片高度 = 112;  // 调整为标准人脸尺寸
    private static final int 图片宽度 = 112;
    private static final int 图片通道数 = 3;
    private static final String 特征数据文件 = "face_features.dat";
    private static final int 特征提取次数 = 3;  // 减少特征提取次数
    private static final int 数据增强次数 = 3;   // 减少数据增强次数
    private static final int 特征维度 = 512;   // 保持特征维度
    private static final int 并行线程数 = 4;
    private static final double 特征融合权重 = 0.9;
    private static final double 余弦相似度权重 = 0.4;
    private static final double 欧氏相似度权重 = 0.3;
    private static final double 差异相似度权重 = 0.2;
    private static final double 曼哈顿相似度权重 = 0.1;
    private static final double 相似度差异阈值 = 0.02;
    private static final double 缓存相似度阈值 = 0.95;
    private static final double 识别相似度阈值 = 0.90;
    private Map<Integer, INDArray> 特征缓存;
    private static final int 缓存大小 = 100;
    private ExecutorService 线程池;

    public FaceRecognitionModel() {
        this.人脸映射表 = new HashMap<>();
        this.人脸特征库 = new HashMap<>();
        this.特征缓存 = new LinkedHashMap<Integer, INDArray>(缓存大小, 0.75f, true) {
            @Override
            protected boolean removeEldestEntry(Map.Entry<Integer, INDArray> eldest) {
                return size() > 缓存大小;
            }
        };
        this.线程池 = Executors.newFixedThreadPool(并行线程数);
        构建模型();
    }

    private void 构建模型() {
        MultiLayerConfiguration 配置 = new NeuralNetConfiguration.Builder()
            .seed(123)
            .weightInit(WeightInit.XAVIER)
            .updater(new Adam(new ExponentialSchedule(ScheduleType.ITERATION, 0.001, 0.95)))
            .l2(0.0001)
            .list()
            // 第一个卷积块
            .layer(new ConvolutionLayer.Builder(3, 3)
                .nIn(图片通道数)
                .stride(1, 1)
                .nOut(32)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            .layer(new DropoutLayer.Builder(0.2).build())

            // 第二个卷积块
            .layer(new ConvolutionLayer.Builder(3, 3)
                .stride(1, 1)
                .nOut(64)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            .layer(new DropoutLayer.Builder(0.2).build())

            // 第三个卷积块
            .layer(new ConvolutionLayer.Builder(3, 3)
                .stride(1, 1)
                .nOut(128)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            .layer(new DropoutLayer.Builder(0.2).build())

            // 第四个卷积块
            .layer(new ConvolutionLayer.Builder(3, 3)
                .stride(1, 1)
                .nOut(256)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            .layer(new DropoutLayer.Builder(0.2).build())

            // 全局池化层
            .layer(new GlobalPoolingLayer.Builder(PoolingType.AVG).build())

            // 全连接层
            .layer(new DenseLayer.Builder()
                .nOut(1024)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new DropoutLayer.Builder(0.5).build())

            .layer(new DenseLayer.Builder()
                .nOut(特征维度)
                .activation(Activation.RELU)
                .build())
            .layer(new BatchNormalization.Builder().build())
            .layer(new DropoutLayer.Builder(0.5).build())

            // 输出层
            .layer(new OutputLayer.Builder(LossFunctions.LossFunction.COSINE_PROXIMITY)
                .nOut(特征维度)
                .activation(Activation.IDENTITY)
                .build())
            .setInputType(InputType.convolutional(图片高度, 图片宽度, 图片通道数))
            .build();

        模型 = new MultiLayerNetwork(配置);
        模型.init();
    }

    private INDArray 数据增强(INDArray 原始图片) {
        if (原始图片.rank() != 4) {
            throw new IllegalArgumentException("输入图片必须是4维张量 [batch, channels, height, width]");
        }

        int 方法 = new Random().nextInt(5);
        switch (方法) {
            case 0:  // 水平翻转
                return 水平翻转(原始图片);
            case 1:  // 随机裁剪
                return 随机裁剪(原始图片, 0.9);
            case 2:  // 添加高斯噪声
                return 添加高斯噪声(原始图片, 0.01);
            case 3:  // 随机遮挡
                return 随机遮挡(原始图片, 0.2);
            case 4:  // 亮度调整
                return 调整亮度(原始图片, 0.9, 1.1);
            default:
                return 原始图片;
        }
    }

    private INDArray 水平翻转(INDArray 图片) {
        INDArray 翻转图片 = 图片.dup();
        for (int c = 0; c < 图片通道数; c++) {
            for (int i = 0; i < 图片高度; i++) {
                for (int j = 0; j < 图片宽度/2; j++) {
                    int 对应列 = 图片宽度 - 1 - j;
                    double 临时值 = 翻转图片.getDouble(new int[]{0, c, i, j});
                    翻转图片.putScalar(new int[]{0, c, i, j}, 翻转图片.getDouble(new int[]{0, c, i, 对应列}));
                    翻转图片.putScalar(new int[]{0, c, i, 对应列}, 临时值);
                }
            }
        }
        return 翻转图片;
    }

    private INDArray 随机裁剪(INDArray 图片, double 比例) {
        int 新高度 = (int)(图片高度 * 比例);
        int 新宽度 = (int)(图片宽度 * 比例);
        int 起始行 = new Random().nextInt(图片高度 - 新高度);
        int 起始列 = new Random().nextInt(图片宽度 - 新宽度);
        return 图片.get(NDArrayIndex.all(), NDArrayIndex.all(),
                        NDArrayIndex.interval(起始行, 起始行 + 新高度),
                        NDArrayIndex.interval(起始列, 起始列 + 新宽度));
    }

    private INDArray 添加高斯噪声(INDArray 图片, double 标准差) {
        INDArray 噪声 = Nd4j.randn(图片.shape()).mul(标准差);
        return 图片.add(噪声);
    }

    private INDArray 随机遮挡(INDArray 图片, double 遮挡比例) {
        INDArray 遮挡图片 = 图片.dup();
        int 遮挡高度 = (int)(图片高度 * 遮挡比例);
        int 遮挡宽度 = (int)(图片宽度 * 遮挡比例);
        int 起始行 = new Random().nextInt(图片高度 - 遮挡高度);
        int 起始列 = new Random().nextInt(图片宽度 - 遮挡宽度);
        遮挡图片.get(NDArrayIndex.all(), NDArrayIndex.all(),
                     NDArrayIndex.interval(起始行, 起始行 + 遮挡高度),
                     NDArrayIndex.interval(起始列, 起始列 + 遮挡宽度)).assign(0);
        return 遮挡图片;
    }

    private INDArray 调整亮度(INDArray 图片, double 最小因子, double 最大因子) {
        double 因子 = 最小因子 + new Random().nextDouble() * (最大因子 - 最小因子);
        return 图片.mul(因子);
    }

    private INDArray 预处理图片(INDArray 原始图片) {
        // 确保输入图片是4维张量 [batch, channels, height, width]
        if (原始图片.rank() != 4) {
            throw new IllegalArgumentException("输入图片必须是4维张量 [batch, channels, height, width]");
        }

        // 检查并调整通道数
        if (原始图片.size(1) != 图片通道数) {
            if (原始图片.size(1) == 1) {
                // 如果是单通道，复制到3个通道
                INDArray 新图片 = Nd4j.zeros(原始图片.size(0), 图片通道数, 原始图片.size(2), 原始图片.size(3));
                for (int i = 0; i < 图片通道数; i++) {
                    新图片.get(NDArrayIndex.all(), NDArrayIndex.point(i), NDArrayIndex.all(), NDArrayIndex.all())
                        .assign(原始图片.get(NDArrayIndex.all(), NDArrayIndex.point(0), NDArrayIndex.all(), NDArrayIndex.all()));
                }
                原始图片 = 新图片;
            } else {
                throw new IllegalArgumentException("输入图片通道数必须是1或3");
            }
        }

        // 调整图片尺寸
        if (原始图片.size(2) != 图片高度 || 原始图片.size(3) != 图片宽度) {
            // 使用双线性插值调整大小
            INDArray 调整后图片 = Nd4j.zeros(原始图片.size(0), 图片通道数, 图片高度, 图片宽度);
            double 高度比例 = (double)原始图片.size(2) / 图片高度;
            double 宽度比例 = (double)原始图片.size(3) / 图片宽度;

            for (int h = 0; h < 图片高度; h++) {
                for (int w = 0; w < 图片宽度; w++) {
                    double 原始h = h * 高度比例;
                    double 原始w = w * 宽度比例;
                    
                    int h1 = (int)Math.floor(原始h);
                    int h2 = Math.min(h1 + 1, (int)原始图片.size(2) - 1);
                    int w1 = (int)Math.floor(原始w);
                    int w2 = Math.min(w1 + 1, (int)原始图片.size(3) - 1);
                    
                    double h_ratio = 原始h - h1;
                    double w_ratio = 原始w - w1;
                    
                    for (int c = 0; c < 图片通道数; c++) {
                        double v1 = 原始图片.getDouble(0, c, h1, w1);
                        double v2 = 原始图片.getDouble(0, c, h1, w2);
                        double v3 = 原始图片.getDouble(0, c, h2, w1);
                        double v4 = 原始图片.getDouble(0, c, h2, w2);
                        
                        double value = (1 - h_ratio) * (1 - w_ratio) * v1 +
                                     (1 - h_ratio) * w_ratio * v2 +
                                     h_ratio * (1 - w_ratio) * v3 +
                                     h_ratio * w_ratio * v4;
                        
                        调整后图片.putScalar(new int[]{0, c, h, w}, value);
                    }
                }
            }
            原始图片 = 调整后图片;
        }

        // 归一化像素值到[0,1]范围
        double 最大值 = 原始图片.maxNumber().doubleValue();
        if (最大值 > 1.0) {
            原始图片.divi(255.0);
        }

        // 确保图片尺寸正确
        if (原始图片.size(2) != 图片高度 || 原始图片.size(3) != 图片宽度) {
            throw new IllegalStateException(String.format(
                "图片尺寸调整失败：期望 [%d, %d]，实际 [%d, %d]",
                图片高度, 图片宽度, 原始图片.size(2), 原始图片.size(3)
            ));
        }

        return 原始图片;
    }

    public void 注册人脸(INDArray 图片数据, int 人员ID, String 姓名) throws IOException {
        final INDArray 预处理后图片 = 预处理图片(图片数据);
        
        List<Future<INDArray>> 特征任务列表 = new ArrayList<>();
        
        // 并行提取特征
        for (int i = 0; i < 特征提取次数; i++) {
            final int 索引 = i;
            Future<INDArray> 任务 = 线程池.submit(() -> {
                if (索引 == 0) {
                    return 模型.output(预处理后图片);
                } else {
                    INDArray 增强图片 = 数据增强(预处理后图片);
                    return 模型.output(增强图片);
                }
            });
            特征任务列表.add(任务);
        }
        
        // 收集特征
        List<INDArray> 特征列表 = new ArrayList<>();
        for (Future<INDArray> 任务 : 特征任务列表) {
            try {
                特征列表.add(任务.get());
            } catch (Exception e) {
                throw new IOException("特征提取失败", e);
            }
        }
        
        // 计算加权平均特征向量
        INDArray 平均特征 = Nd4j.zeros(特征列表.get(0).shape());
        double 总权重 = 0;
        for (int i = 0; i < 特征列表.size(); i++) {
            double 权重 = i == 0 ? 1.0 : 特征融合权重;
            平均特征.addi(特征列表.get(i).mul(权重));
            总权重 += 权重;
        }
        平均特征.divi(总权重);
        
        // 特征归一化
        平均特征 = 特征归一化(平均特征);
        
        人脸特征库.put(人员ID, Collections.singletonList(平均特征));
        人脸映射表.put(人员ID, 姓名);
        保存特征数据();
    }

    public void 保存模型(String 模型路径) throws IOException {
        File 模型文件 = new File(模型路径);
        ModelSerializer.writeModel(模型, 模型文件, true);
    }

    public void 加载模型(String 模型路径) throws IOException {
        File 模型文件 = new File(模型路径);
        模型 = ModelSerializer.restoreMultiLayerNetwork(模型文件);
        // 加载特征数据
        加载特征数据();
    }

    private void 保存特征数据() throws IOException {
        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(特征数据文件))) {
            // 保存人脸映射表
            oos.writeObject(人脸映射表);
            
            // 保存特征库
            Map<Integer, List<double[]>> 序列化特征库 = new HashMap<>();
            for (Map.Entry<Integer, List<INDArray>> entry : 人脸特征库.entrySet()) {
                List<double[]> 序列化特征列表 = new ArrayList<>();
                for (INDArray 特征 : entry.getValue()) {
                    序列化特征列表.add(特征.toDoubleVector());
                }
                序列化特征库.put(entry.getKey(), 序列化特征列表);
            }
            oos.writeObject(序列化特征库);
        }
    }

    @SuppressWarnings("unchecked")
    private void 加载特征数据() throws IOException {
        File 特征文件 = new File(特征数据文件);
        if (!特征文件.exists()) {
            return;
        }

        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(特征文件))) {
            // 加载人脸映射表
            人脸映射表 = (Map<Integer, String>) ois.readObject();
            
            // 加载特征库
            Map<Integer, List<double[]>> 序列化特征库 = (Map<Integer, List<double[]>>) ois.readObject();
            人脸特征库.clear();
            for (Map.Entry<Integer, List<double[]>> entry : 序列化特征库.entrySet()) {
                List<INDArray> 特征列表 = new ArrayList<>();
                for (double[] 特征数据 : entry.getValue()) {
                    特征列表.add(Nd4j.create(特征数据));
                }
                人脸特征库.put(entry.getKey(), 特征列表);
            }
        } catch (ClassNotFoundException e) {
            throw new IOException("加载特征数据失败", e);
        }
    }

    public Map<String, Object> 识别人脸(INDArray 图片数据, double 相似度阈值) {
        final INDArray 预处理后图片 = 预处理图片(图片数据);
        
        // 检查缓存
        Integer 缓存ID = 检查缓存(预处理后图片);
        if (缓存ID != null) {
            return 创建结果(缓存ID, 人脸映射表.get(缓存ID), 1.0);
        }

        // 提取特征
        final INDArray 特征 = 特征归一化(模型.output(预处理后图片));
        
        // 并行计算相似度
        List<Future<Map.Entry<Integer, Double>>> 相似度任务列表 = new ArrayList<>();
        for (Map.Entry<Integer, List<INDArray>> entry : 人脸特征库.entrySet()) {
            final int ID = entry.getKey();
            final INDArray 注册特征 = entry.getValue().get(0);
            Future<Map.Entry<Integer, Double>> 任务 = 线程池.submit(() -> {
                double 相似度 = 计算改进相似度(特征, 注册特征);
                return new AbstractMap.SimpleEntry<>(ID, 相似度);
            });
            相似度任务列表.add(任务);
        }
        
        // 收集结果
        double 最大相似度 = -1;
        int 最匹配ID = -1;
        double 次大相似度 = -1;
        
        for (Future<Map.Entry<Integer, Double>> 任务 : 相似度任务列表) {
            try {
                Map.Entry<Integer, Double> 结果 = 任务.get();
                if (结果.getValue() > 最大相似度) {
                    次大相似度 = 最大相似度;
                    最大相似度 = 结果.getValue();
                    最匹配ID = 结果.getKey();
                } else if (结果.getValue() > 次大相似度) {
                    次大相似度 = 结果.getValue();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        // 如果最大相似度超过阈值，直接返回结果
        if (最大相似度 >= 识别相似度阈值) {
            // 更新缓存
            特征缓存.put(最匹配ID, 特征);
            return 创建结果(最匹配ID, 人脸映射表.get(最匹配ID), 最大相似度);
        }

        // 如果最大相似度和次大相似度太接近，可能是不确定的情况
        if (最大相似度 - 次大相似度 < 相似度差异阈值) {
            return 创建结果(-1, "未知", 最大相似度);
        }

        // 如果最大相似度在阈值范围内，返回结果
        if (最大相似度 >= 相似度阈值) {
            // 更新缓存
            特征缓存.put(最匹配ID, 特征);
            return 创建结果(最匹配ID, 人脸映射表.get(最匹配ID), 最大相似度);
        }

        return 创建结果(-1, "未知", 最大相似度);
    }

    private Integer 检查缓存(INDArray 图片数据) {
        INDArray 特征 = 特征归一化(模型.output(图片数据));
        for (Map.Entry<Integer, INDArray> entry : 特征缓存.entrySet()) {
            if (计算改进相似度(特征, entry.getValue()) > 缓存相似度阈值) {
                return entry.getKey();
            }
        }
        return null;
    }

    private Map<String, Object> 创建结果(int ID, String 姓名, double 相似度) {
        Map<String, Object> 结果 = new HashMap<>();
        结果.put("ID", ID);
        结果.put("姓名", 姓名);
        结果.put("相似度", 相似度);
        return 结果;
    }

    private double 计算改进相似度(INDArray 特征1, INDArray 特征2) {
        特征1 = 特征归一化(特征1);
        特征2 = 特征归一化(特征2);
        
        // 计算余弦相似度
        double 余弦相似度 = 特征1.mul(特征2).sumNumber().doubleValue();
        
        // 计算欧氏距离
        double 欧氏距离 = Math.sqrt(特征1.sub(特征2).mul(特征1.sub(特征2)).sumNumber().doubleValue());
        double 欧氏相似度 = 1.0 / (1.0 + 欧氏距离);
        
        // 计算特征差异
        INDArray 差异 = 特征1.sub(特征2);
        double 差异范数 = Math.sqrt(差异.mul(差异).sumNumber().doubleValue());
        double 差异相似度 = Math.exp(-差异范数);
        
        // 计算曼哈顿距离
        double 曼哈顿距离 = 差异.norm1Number().doubleValue();
        double 曼哈顿相似度 = 1.0 / (1.0 + 曼哈顿距离);
        
        // 融合四种相似度
        return 余弦相似度权重 * 余弦相似度 + 
               欧氏相似度权重 * 欧氏相似度 + 
               差异相似度权重 * 差异相似度 +
               曼哈顿相似度权重 * 曼哈顿相似度;
    }

    private INDArray 特征归一化(INDArray 特征) {
        // 确保特征向量不为零
        if (特征.sumNumber().doubleValue() == 0) {
            return 特征;
        }
        double 范数 = Math.sqrt(特征.mul(特征).sumNumber().doubleValue());
        return 特征.div(范数);
    }

    public void 关闭() {
        线程池.shutdown();
        try {
            if (!线程池.awaitTermination(60, TimeUnit.SECONDS)) {
                线程池.shutdownNow();
            }
        } catch (InterruptedException e) {
            线程池.shutdownNow();
        }
        模型.close();
    }
} 